import numpy as np
import torch
import torch.nn as nn
from src.diff_models_table import diff_CSDI
import yaml
import torch.nn.functional as F  # 添加这行导入


class CSDI_base(nn.Module):
    def __init__(self, target_dim, config, device, max_values=None, min_values=None,corr_rules=None):
        super().__init__()
        self.device = device
        self.target_dim = target_dim

        self.emb_time_dim = config["model"]["timeemb"]
        self.emb_feature_dim = config["model"]["featureemb"]

        self.is_unconditional = config["model"]["is_unconditional"]
        self.target_strategy = config["model"]["target_strategy"]

        self.emb_total_dim = self.emb_time_dim + self.emb_feature_dim
        if self.is_unconditional == False:
            self.emb_total_dim += 1  # for conditional mask

        self.embed_layer = nn.Embedding(
            num_embeddings=self.target_dim, embedding_dim=self.emb_feature_dim
        )

        config_diff = config["diffusion"]
        config_diff["side_dim"] = self.emb_total_dim

        # 新增相关性约束参数
        self.corr_rules = corr_rules  # 从dataset传入的规则
        # self.corr_loss_weight = 0.5  # 可调整的权重系数

        input_dim = 1 if self.is_unconditional == True else 2
        self.diffmodel = diff_CSDI(config_diff, input_dim)

        # parameters for diffusion models
        self.num_steps = config_diff["num_steps"]
        if config_diff["schedule"] == "quad":
            self.beta = (
                np.linspace(
                    config_diff["beta_start"] ** 0.5,
                    config_diff["beta_end"] ** 0.5,
                    self.num_steps,
                )
                ** 2
            )
        elif config_diff["schedule"] == "linear":
            self.beta = np.linspace(
                config_diff["beta_start"], config_diff["beta_end"], self.num_steps
            )

        self.alpha_hat = 1 - self.beta
        self.alpha = np.cumprod(self.alpha_hat)
        self.alpha_torch = (
            torch.tensor(self.alpha).float().to(self.device).unsqueeze(1).unsqueeze(1)
        )

        # 新增最大值和最小值约束
        self.max_values = torch.tensor(max_values * 1.1, device=self.device)  # 放宽 10%
        self.min_values = torch.tensor(min_values * 0.9, device=self.device)  # 放宽 10%

        # 如果 max_values 和 min_values 是 numpy 数组，将其转换为 tensor
        # 修改极值张量初始化部分
        if max_values is not None:
            if isinstance(max_values, np.ndarray):
                max_values = torch.from_numpy(max_values).float()
            self.max_values = max_values.view(1, -1, 1).to(device) * 1.1  # [1, num_features, 1]

        if min_values is not None:
            if isinstance(min_values, np.ndarray):
                min_values = torch.from_numpy(min_values).float()
            self.min_values = min_values.view(1, -1, 1).to(device) * 0.9  # [1, num_features, 1]

    def time_embedding(self, pos, d_model=128):
        pe = torch.zeros(pos.shape[0], pos.shape[1], d_model).to(self.device)
        position = pos.unsqueeze(2)
        div_term = 1 / torch.pow(
            10000.0, torch.arange(0, d_model, 2).to(self.device) / d_model
        )
        pe[:, :, 0::2] = torch.sin(position * div_term)
        pe[:, :, 1::2] = torch.cos(position * div_term)
        return pe

    def get_randmask(self, observed_mask):
        rand_for_mask = torch.rand_like(observed_mask) * observed_mask
        rand_for_mask = rand_for_mask.reshape(len(rand_for_mask), -1)

        for i in range(len(observed_mask)):
            sample_ratio = 0.8  # np.random.rand()  # missing ratio
            num_observed = observed_mask[i].sum().item()
            num_masked = round(num_observed * sample_ratio)
            rand_for_mask[i][rand_for_mask[i].topk(num_masked).indices] = -1
        cond_mask = (rand_for_mask > 0).reshape(observed_mask.shape).float()
        return cond_mask

    def get_side_info(self, observed_tp, cond_mask):
        B, K, L = cond_mask.shape

        time_embed = self.time_embedding(observed_tp, self.emb_time_dim)  # (B,L,emb)
        time_embed = time_embed.unsqueeze(2).expand(-1, -1, K, -1)
        feature_embed = self.embed_layer(
            torch.arange(self.target_dim).to(self.device)
        )  # (K,emb)
        feature_embed = feature_embed.unsqueeze(0).unsqueeze(0).expand(B, L, -1, -1)

        side_info = torch.cat([time_embed, feature_embed], dim=-1)  # (B,L,K,*)
        side_info = side_info.permute(0, 3, 2, 1)  # (B,*,K,L)

        if self.is_unconditional == False:
            side_mask = cond_mask.unsqueeze(1)  # (B,1,K,L)
            side_info = torch.cat([side_info, side_mask], dim=1)

        return side_info

    def calc_loss_valid(
        self, observed_data, cond_mask, observed_mask, side_info, is_train
    ):
        loss_sum = 0
        # In validation, perform T steps forward and backward.
        for t in range(self.num_steps):
            loss = self.calc_loss(
                observed_data, cond_mask, observed_mask, side_info, is_train, set_t=t
            )
            loss_sum += loss.detach()
        return loss_sum / self.num_steps

    def _apply_correlation_constraint(self, predicted, observed_mask):
        """计算相关性约束损失"""
        if self.corr_rules is None or len(self.corr_rules) == 0:
            return torch.tensor(0.0, device=self.device)

        # 确保预测值的维度是 (batch_size, num_features, seq_len)
        if predicted.dim() == 2:
            predicted = predicted.unsqueeze(1)  # [B, L] -> [B, 1, L]

        batch_size, num_features, seq_len = predicted.shape

        loss = torch.tensor(0.0, device=self.device)
        valid_rules = 0

        for rule in self.corr_rules:
            i, j = rule['feat_pair']

            # 检查特征索引是否有效
            if i >= num_features or j >= num_features:
                print(f"Warning: Correlation rule {i}-{j} skipped (num_features={num_features})")
                continue

            # 获取特征对 [B, 2, L]
            pred_pair = predicted[:, [i, j], :]

            # 展平为 [2, B*L] 用于计算相关性
            pred_flatten = pred_pair.permute(1, 0, 2).reshape(2, -1)

            # 计算皮尔逊相关系数
            if pred_flatten.shape[1] < 2:  # 需要至少2个点计算相关性
                continue

            actual_corr = torch.corrcoef(pred_flatten)[0, 1]

            # 计算Huber损失
            target_corr = torch.tensor(1.0 if rule['corr_type'] == 1 else -1.0,
                                       device=self.device)
            loss += F.huber_loss(actual_corr, target_corr, delta=0.2)
            valid_rules += 1

        # 返回平均损失（仅对有效规则）
        return loss / valid_rules if valid_rules > 0 else torch.tensor(0.0, device=self.device)

    # 计算损失
    def calc_loss(self, observed_data, cond_mask, observed_mask, side_info, is_train, set_t=-1):
        # 确保输入维度正确 [B, K, L]
        if observed_data.dim() == 2:
            observed_data = observed_data.unsqueeze(1)  # [B, L] -> [B, 1, L]
        B, K, L = observed_data.shape

        # 时间步采样
        if is_train != 1:  # validation
            t = (torch.ones(B) * set_t).long().to(self.device)
        else:
            t = torch.randint(0, self.num_steps, [B]).to(self.device)

        # 扩散过程
        current_alpha = self.alpha_torch[t].view(-1, 1, 1)  # (B,1,1)
        noise = torch.randn_like(observed_data)
        noisy_data = (current_alpha ** 0.5) * observed_data + (1.0 - current_alpha) ** 0.5 * noise

        # 模型预测
        total_input = self.set_input_to_diffmodel(noisy_data, observed_data, cond_mask)
        predicted = self.diffmodel(total_input, side_info, t)  # (B,K,L)

        # 基础MSE损失
        target_mask = observed_mask - cond_mask
        residual = (noise - predicted) * target_mask
        num_eval = target_mask.sum()
        mse_loss = (residual ** 2).sum() / (num_eval if num_eval > 0 else 1)

        # 初始化总损失
        total_loss = mse_loss

        # 1. 值范围约束 (仅当极值存在时计算)
        if hasattr(self, 'max_values') and hasattr(self, 'min_values'):
            # 确保极值张量维度 [1, K, 1] 与预测值 [B, K, L] 兼容
            max_vals = self.max_values.expand(-1, -1, L)  # [1, K, 1] -> [1, K, L]
            min_vals = self.min_values.expand(-1, -1, L)

            # 仅对目标区域计算约束
            valid_pred = predicted * target_mask
            valid_max = max_vals * target_mask
            valid_min = min_vals * target_mask

            # 计算约束损失（改进版）
            max_constraint = torch.where(
                valid_pred > valid_max,
                F.huber_loss(valid_pred, valid_max, reduction='none', delta=0.1),
                torch.zeros_like(valid_pred)
            )

            min_constraint = torch.where(
                valid_pred < valid_min,
                F.huber_loss(valid_pred, valid_min, reduction='none', delta=0.1),
                torch.zeros_like(valid_pred)
            )

            constraint_loss = (max_constraint + min_constraint).sum() / (num_eval + 1e-6)
            total_loss = total_loss + 0.5 * constraint_loss

        # 2. 相关性约束（保持不变）
        if hasattr(self, 'corr_rules') and self.corr_rules:
            corr_loss = self._apply_correlation_constraint(
                predicted * target_mask,
                observed_mask
            )
            total_loss = total_loss + 0.5 * corr_loss

        return total_loss

    def set_input_to_diffmodel(self, noisy_data, observed_data, cond_mask):
        if self.is_unconditional == True:
            total_input = noisy_data.unsqueeze(1)  # (B,1,K,L)
        else:
            cond_obs = (cond_mask * observed_data).unsqueeze(1)
            noisy_target = ((1 - cond_mask) * noisy_data).unsqueeze(1)
            total_input = torch.cat([cond_obs, noisy_target], dim=1)  # (B,2,K,L)

        return total_input

    def impute(self, observed_data, cond_mask, side_info, n_samples, noisy_cond_history=None):
        B, K, L = observed_data.shape
        imputed_samples = torch.zeros(B, n_samples, K, L).to(self.device)

        for i in range(n_samples):
            current_sample = torch.randn_like(observed_data)

            for t in range(self.num_steps - 1, -1, -1):
                if self.is_unconditional == True:
                    diff_input = (
                            cond_mask * noisy_cond_history[t]
                            + (1.0 - cond_mask) * current_sample
                    )
                    diff_input = diff_input.unsqueeze(1)  # (B,1,K,L)
                else:
                    cond_obs = (cond_mask * observed_data).unsqueeze(1)
                    noisy_target = ((1 - cond_mask) * current_sample).unsqueeze(1)
                    diff_input = torch.cat([cond_obs, noisy_target], dim=1)  # (B,2,K,L)
                predicted = self.diffmodel(
                    diff_input, side_info, torch.tensor([t]).to(self.device)
                )  # (B,K,L)

                coeff1 = 1 / self.alpha_hat[t] ** 0.5
                coeff2 = (1 - self.alpha_hat[t]) / (1 - self.alpha[t]) ** 0.5

                current_sample = coeff1 * (current_sample - coeff2 * predicted)

                if t > 0:
                    noise = torch.randn_like(current_sample)
                    sigma = (
                                    (1.0 - self.alpha[t - 1]) / (1.0 - self.alpha[t]) * self.beta[t]
                            ) ** 0.5
                    current_sample += sigma * noise

            # 硬性约束 截断
            if self.max_values is not None and self.min_values is not None:
                current_sample = torch.clamp(current_sample, min=self.min_values, max=self.max_values)

            imputed_samples[:, i] = current_sample.detach()
        return imputed_samples

    def forward(self, batch, is_train=1):
        # 确保所有输入张量都在正确设备上
        batch = {k: v.to(self.device) if torch.is_tensor(v) else v
                 for k, v in batch.items()}
        (
            observed_data,
            observed_mask,
            observed_tp,
            gt_mask,
            for_pattern_mask,
            _,
        ) = self.process_data(batch)

        # In testing, using `gt_mask` (generated with fixed missing rate).
        if is_train == 0:
            cond_mask = gt_mask
        # In training, generate random mask
        else:
            cond_mask = self.get_randmask(observed_mask)

        side_info = self.get_side_info(observed_tp, cond_mask)

        loss_func = self.calc_loss if is_train == 1 else self.calc_loss_valid
        return loss_func(observed_data, cond_mask, observed_mask, side_info, is_train)

    def evaluate(self, batch, n_samples):
        (
            observed_data,
            observed_mask,
            observed_tp,
            gt_mask,
            _,
            cut_length,
        ) = self.process_data(batch)

        with torch.no_grad():
            cond_mask = gt_mask
            target_mask = observed_mask - cond_mask
            side_info = self.get_side_info(observed_tp, cond_mask)

            samples = self.impute(observed_data, cond_mask, side_info, n_samples)
            print(len(samples))
            print(len(observed_data))

        return samples, observed_data, target_mask, observed_mask, observed_tp


class TabCSDI(CSDI_base):
    def __init__(self, config, device, target_dim=1, max_values=None, min_values=None, corr_rules=None):
        super().__init__(target_dim, config, device, max_values=max_values, min_values=min_values, corr_rules=corr_rules)

    def process_data(self, batch):
        """处理输入数据，确保维度一致性和设备正确性

        Args:
            batch: 包含以下键的字典:
                - observed_data: 观测数据 [B, L]
                - observed_mask: 观测掩码 [B, L]
                - timepoints: 时间点 [B, L] 或 [L]
                - gt_mask: 真实值掩码 [B, L]

        Returns:
            Tuple: (observed_data, observed_mask, observed_tp, gt_mask, for_pattern_mask, cut_length)
        """
        # 1. 观测数据 [B, L] -> [B, 1, L]
        observed_data = batch["observed_data"]
        if observed_data.dim() == 2 and observed_data.shape[1] > 1:  # 确保是特征维度扩展
            observed_data = observed_data.unsqueeze(1)  # [B, L] -> [B, 1, L]
        observed_data = observed_data.float().to(self.device)

        # 2. 观测掩码 [B, L] -> [B, 1, L]
        observed_mask = batch["observed_mask"]
        if observed_mask.dim() == 2:
            observed_mask = observed_mask.unsqueeze(1)
        observed_mask = observed_mask.float().to(self.device)

        # 3. 时间点处理 (兼容不同输入维度)
        observed_tp = batch["timepoints"]
        if observed_tp.dim() == 1:  # [L] -> [B, L]
            observed_tp = observed_tp.unsqueeze(0).expand(observed_data.shape[0], -1)
        observed_tp = observed_tp.float().to(self.device)

        # 4. 真实值掩码 [B, L] -> [B, 1, L]
        gt_mask = batch["gt_mask"]
        if gt_mask.dim() == 2:
            gt_mask = gt_mask.unsqueeze(1)
        gt_mask = gt_mask.float().to(self.device)

        # 5. 辅助变量初始化
        cut_length = torch.zeros(observed_data.size(0), dtype=torch.long, device=self.device)
        for_pattern_mask = observed_mask.clone()  # 深拷贝避免意外修改

        # 维度验证
        assert observed_data.dim() == 3, f"观测数据应为3维，实际得到{observed_data.shape}"
        assert observed_mask.shape == observed_data.shape, "掩码与数据形状不匹配"

        return (
            observed_data,  # [B, 1, L]
            observed_mask,  # [B, 1, L]
            observed_tp,  # [B, L]
            gt_mask,  # [B, 1, L]
            for_pattern_mask,  # [B, 1, L]
            cut_length,  # [B]
        )